/*
 * CatGraph.cpp
 *
 *  Created on: 10/11/2011
 *      Author: will
 */

#include "CatGraph.h"

#define VIS_DIV 100
#define NEW_MUL 50

static const double eps = 2.22044604925031e-16;

CatGraph::CatGraph(unsigned int numParticles, unsigned int numNodes,
		Matrix3d Q, Vector3d u_mean) {
	frameNum = 0;
	this->numParticles = numParticles;
	this->numNodes = numNodes;
	this->Q = Q;
	this->u_mean = u_mean;
	particles.clear();
	particles.reserve(numParticles);
	nodes.clear();
	nodes.reserve(numNodes);

	information.clear();
	informationLookup.clear();
	updateList.clear();

	codebook.release();
	matcher = NULL;
}

void CatGraph::initialiseVision(string codebookPath, string chowLiuPath) {
	codebook.release();
	cv::FileStorage codebookFile(codebookPath,cv::FileStorage::READ);
    codebookFile["codebook"] >> codebook;
    codebookFile.release();

    // Memory leaky...
    matcher = new cv::FlannBasedMatcher();

    clTree.load(chowLiuPath);

    double PZGE = 0.39, PZGNE = 0.0;
    fabMap = FastLookupFabMap(&clTree,PZGE,PZGNE);
    z_avg.setAsAvgPlace(&clTree,-1,PZGE,PZGNE);

    mexPrintf("Codebook with %d %d-D words loaded.\n",codebook.rows,codebook.cols);

}

void CatGraph::initialiseTrajectory(Mat* descriptors, string descriptorPath) {

	vector<cv::DMatch> matches;
	Bagofwords* bagOfWords = new Bagofwords(codebook.rows);
	matcher->match(*descriptors,codebook,matches);
	for (vector<cv::DMatch>::iterator matchIter = matches.begin(); matchIter
			!= matches.end(); matchIter++) {
		bagOfWords->add(matchIter->trainIdx);
	}

	GraphNode* node = new GraphNode(frameNum, descriptors, bagOfWords, descriptorPath);

	//GraphNode* node = new GraphNode(frameNum, descriptors, NULL);
	nodes.push_back(node);
	frameNum++;
}

void CatGraph::initialiseParticles(unsigned int blankFrames) {
	Particle* particle;
	for (unsigned int i = 0; i < numParticles; i++) {
		particle = randomParticle(blankFrames);
		particle->weight = 1 / (double) numParticles;
		particles.push_back(particle);
	}
}

void CatGraph::addNode(Mat* descriptors, string descriptorPath, Vector3d odometry) {


	vector<cv::DMatch> matches;
	Bagofwords* bagOfWords = new Bagofwords(codebook.rows);
	matcher->match(*descriptors,codebook,matches);
	for (vector<cv::DMatch>::iterator matchIter = matches.begin(); matchIter
			!= matches.end(); matchIter++) {
		bagOfWords->add(matchIter->trainIdx);
	}

	GraphNode* node = new GraphNode(frameNum, descriptors, bagOfWords, descriptorPath);

	//GraphNode* node = new GraphNode(frameNum, descriptors, NULL);
	Matrix3d icov = getInverseCovariance(odometry, Q);
	GraphEdge* edge = nodes.back()->addEdgeTo(node, odometry, icov);

	updateList.push_back(nodes.back());
	nodes.push_back(node);
	edges.push_back(edge);

	mexPrintf("Added node %X at %d with %d words to graph.\n", node,
			node->frameNum, node->descriptors->rows);
	frameNum++;

	// For FAB-MAP only - clear descriptors to avoid massive memory reqs
	descriptors->release();
}


void CatGraph::updateInformation() {

	// Update list based method

	mexPrintf("Information Update list has %d nodes\n", updateList.size());

	GraphNode* node;
	for (list<GraphNode*>::iterator nodeIter = updateList.begin(); nodeIter
			!= updateList.end(); nodeIter++) {

	// Full calculation method
	//information.clear();
	//updateList.clear();

	//GraphNode* node;
	//for (vector<GraphNode*>::iterator nodeIter = nodes.begin(); nodeIter
	//			!= nodes.end(); nodeIter++) {

		node = *nodeIter;

		mexPrintf("Information Update: Node %X at %d\n",node,node->frameNum);
		if (node->edgeFrom.size() && node->edgeTo.size()) {

			// build local map to depth of 1
			LocalGraph * localGraph = new LocalGraph(node, (unsigned int) 1);

			node->odoLikelihood = -DBL_MAX;
			node->visLikelihood = -DBL_MAX;

			GraphEdge *fromEdge, *toEdge;
			double odoFraction;
			double odoLikelihood;
			double visLikelihood;
			double information, maxInformation = -DBL_MAX;

			for (vector<GraphEdge*>::iterator fromIter = node->edgeFrom.begin(); fromIter
					!= node->edgeFrom.end(); fromIter++) {
				for (vector<GraphEdge*>::iterator toIter = node->edgeTo.begin(); toIter
						!= node->edgeTo.end(); toIter++) {
					fromEdge = *fromIter;
					toEdge = *toIter;

					mexPrintf(
							"Testing link from node %X at %d to node %X at %d\n",
							fromEdge->nodeFrom, fromEdge->nodeFrom->frameNum,
							toEdge->nodeTo, toEdge->nodeTo->frameNum);

					Vector3d origin(0, 0, 0);
					Vector3d odometry = inverseModel(localGraph->nodes.find(
							toEdge->nodeTo)->second, localGraph->nodes.find(
							fromEdge->nodeFrom)->second);
					if (getRadius(odometry) < eps) {
						odometry(1) = odometry(1) + 1e-5;
					}

					mexPrintf("Equivalent odometry %0.4f %0.4f %0.4f\n",
							odometry(0), odometry(1), odometry(2));

					odoFraction = get2DForwardDPx_2(localGraph->nodes.find(
							fromEdge->nodeFrom)->second,
							localGraph->nodes.find(toEdge->nodeTo)->second,
							origin, fromEdge->icov);
					if (odoFraction != odoFraction)
						odoFraction = 0;
					if (odoFraction < 0)
						odoFraction = 0;
					if (odoFraction > 1)
						odoFraction = 1;
					mexPrintf("Calculated odometry fraction %0.4f\n",odoFraction);

					odoLikelihood = get2DForwardLogPx(forwardModelFrac(
							localGraph->nodes.find(fromEdge->nodeFrom)->second,
							odometry, odoFraction), origin, fromEdge->icov);

					mexPrintf("Calculated odometry likelihood %0.4f\n",odoLikelihood);

					visLikelihood = (1 - odoFraction) * fabMap.LLH(*(fromEdge->nodeFrom->bagOfWords),*(node->bagOfWords))
							+ odoFraction * fabMap.LLH(*(toEdge->nodeTo->bagOfWords),*(node->bagOfWords));

					mexPrintf("Calculated visual likelihood %0.4f\n",visLikelihood);

					information = odoLikelihood + visLikelihood / VIS_DIV;

					mexPrintf("Calculated information %0.4f\n",information);

					mexPrintf(
							"  Traversal from Node %X at %d to Node %X at %d:\n    Gives fraction %0.4f odo %0.4f vis %0.4f information %0.4f\n",
							fromEdge->nodeFrom, fromEdge->nodeFrom->frameNum,
							toEdge->nodeTo, toEdge->nodeTo->frameNum,
							odoFraction, odoLikelihood, visLikelihood,
							information);

					if (information > maxInformation) {
						node->odoLikelihood = odoLikelihood;
						node->visLikelihood = visLikelihood;
						node->bestParents = pair<GraphEdge*, GraphEdge*> (
								fromEdge, toEdge);
						maxInformation = information;
					}
				}
			}

			mexPrintf(
				"Best traversal from Node %X at %d to Node %X at %d gives information %0.4f\n",
				node->bestParents.first->nodeFrom, node->bestParents.first->nodeFrom->frameNum,
				node->bestParents.second->nodeTo, node->bestParents.second->nodeTo->frameNum,
				maxInformation);

			if (maxInformation < -DBL_MAX / 2) {
				mexPrintf(
						"!!!! --- Max Information of node %X at %d is very low at %0.4f --- !!!!\n",
						node, node->frameNum, maxInformation);
			} else {


				map<GraphNode*, multimap<double, GraphNode*>::iterator >::iterator
						currInformation = informationLookup.find(node);

				mexPrintf("Found node at %X, end of map at %X\n",currInformation,informationLookup.end());

				if (currInformation != informationLookup.end()) {
					// Value already in map
					mexPrintf(
							"Removing reference from node %X at %d to information %0.4f for node %X at %d\n",
							currInformation->first,
							currInformation->first->frameNum,
							currInformation->second->first,
							currInformation->second->second,
							currInformation->second->second->frameNum);
					this->information.erase(currInformation->second);


					map<double, GraphNode*>::iterator ret =
							this->information.insert(pair<double, GraphNode*> (
									maxInformation, node));
					currInformation->second = ret;

					mexPrintf(
							"Adding reference from node %X at %d to information %0.4f for node %X at %d\n",
							currInformation->first,
							currInformation->first->frameNum,
							currInformation->second->first,
							currInformation->second->second,
							currInformation->second->second->frameNum);

				} else {
					// Value not already in map

					map<double, GraphNode*>::iterator ret =
							this->information.insert(pair<double, GraphNode*> (
									maxInformation, node));
					informationLookup[node] = ret;

					mexPrintf(
							"Adding information %0.4f for node %X at %d\n",
							ret->first,
							ret->second,
							ret->second->frameNum);

				}

				//this->information.insert(pair<double, GraphNode*> (
				//		maxInformation, node));
			}
		} else {
			mexPrintf("Node has insufficient edges, information not calculated\n");

			map<GraphNode*, multimap<double, GraphNode*>::iterator >::iterator
					currInformation = informationLookup.find(node);

			mexPrintf("Found node at %X, end of map at %X\n",currInformation,informationLookup.end());

			if (currInformation != informationLookup.end()) {
				// Value already in map
				mexPrintf(
						"Removing reference from node %X at %d to information %0.4f for node %X at %d\n",
						currInformation->first,
						currInformation->first->frameNum,
						currInformation->second->first,
						currInformation->second->second,
						currInformation->second->second->frameNum);
				this->information.erase(currInformation->second);
				informationLookup.erase(currInformation);
			}
			//node->odoLikelihood = 0;
			//node->visLikelihood = -DBL_MAX;
			//node->bestParents = pair<GraphEdge*, GraphEdge*> (NULL, NULL);
			//
			//catGraph->information.insert(pair<double, GraphNode*> (-DBL_MAX,
			//		node));
		}

	}

	updateList.clear();

}


void CatGraph::pruneNode(GraphNode* node) {

	mexPrintf("Pruning node %X at %d with parents %X and %X\n", node,
			node->frameNum, node->bestParents.first->nodeFrom, node->bestParents.second->nodeTo);
	mexPrintf("Parents %X at %d and %X at %d\n", node->bestParents.first->nodeFrom,
			node->bestParents.first->nodeFrom->frameNum, node->bestParents.second->nodeTo,
			node->bestParents.second->nodeTo->frameNum);

	// build local map to depth of 1
	mexPrintf("build local map to depth of 1\n");
	LocalGraph * localGraph = new LocalGraph(node, (unsigned int) 1);

	// find and remove node and edges from global list
	mexPrintf("find and remove node and edges from global list\n");
	for (vector<GraphNode*>::iterator nodeIter = nodes.begin(); nodeIter
			!= nodes.end(); nodeIter++) {
		if (*nodeIter == node) {
			nodes.erase(nodeIter);
			break;
		}
	}
	for (vector<GraphEdge*>::iterator edgeIter = edges.begin(); edgeIter
			!= edges.end(); edgeIter++) {
		if (localGraph->edges.find(*edgeIter) != localGraph->edges.end()) {
			edges.erase(edgeIter);
			edgeIter--;
		}
	}


	// find and remove node from information content
	mexPrintf("find and remove node from information content\n");
	map<GraphNode*, multimap<double, GraphNode*>::iterator >::iterator
			currInformation = informationLookup.find(node);
	if (currInformation != informationLookup.end()) {
		mexPrintf(
				"Removing reference from node %X at %d to information %0.4f for node %X at %d\n",
				currInformation->first, currInformation->first->frameNum,
				currInformation->second->first,
				currInformation->second->second,
				currInformation->second->second->frameNum);
		information.erase(currInformation->second);
		informationLookup.erase(currInformation);
	} else {
		// BIG ERROR - Tried to remove node without valid information content!
		mexPrintf("BIG ERROR - Tried to remove node without valid information content!\n");
	}

	// add neighbouring nodes to information content recalculation
	mexPrintf("add neighbouring nodes to information content recalculation\n");
	for (map<GraphNode*, Vector3d>::iterator nodeIter =
			localGraph->nodes.begin(); nodeIter != localGraph->nodes.end(); nodeIter++) {
		if (nodeIter->first != node)
			updateList.push_back(nodeIter->first);
	}

	// find and remove edge references from linked nodes
	mexPrintf("find and remove edge references from linked nodes\n");
	for (vector<GraphEdge*>::iterator edgeFromIter = node->edgeFrom.begin(); edgeFromIter
			!= node->edgeFrom.end(); edgeFromIter++) {
		mexPrintf("EdgeFrom %X from %X to %X\n",*edgeFromIter,(*edgeFromIter)->nodeFrom,(*edgeFromIter)->nodeTo);
		for (vector<GraphEdge*>::iterator edgeToIter =
				(*edgeFromIter)->nodeFrom->edgeTo.begin(); edgeToIter
				!= (*edgeFromIter)->nodeFrom->edgeTo.end(); edgeToIter++) {
			mexPrintf("EdgeTo %X from %X to %X\n",*edgeToIter,(*edgeToIter)->nodeFrom,(*edgeToIter)->nodeTo);
			if (*edgeToIter == *edgeFromIter) {
				mexPrintf("  Erasing edge...\n");
				(*edgeFromIter)->nodeFrom->edgeTo.erase(edgeToIter);
				break;
			}
		}
	}
	// find and remove edge references to linked nodes
	mexPrintf("find and remove edge references to linked nodes\n");
	for (vector<GraphEdge*>::iterator edgeToIter = node->edgeTo.begin(); edgeToIter
			!= node->edgeTo.end(); edgeToIter++) {
		mexPrintf("EdgeTo %X from %X to %X\n",*edgeToIter,(*edgeToIter)->nodeFrom,(*edgeToIter)->nodeTo);
		for (vector<GraphEdge*>::iterator edgeFromIter =
				(*edgeToIter)->nodeTo->edgeFrom.begin(); edgeFromIter
				!= (*edgeToIter)->nodeTo->edgeFrom.end(); edgeFromIter++) {
			mexPrintf("EdgeFrom %X from %X to %X\n",*edgeFromIter,(*edgeFromIter)->nodeFrom,(*edgeFromIter)->nodeTo);
			if (*edgeFromIter == *edgeToIter) {
				mexPrintf("  Erasing edge...\n");
				(*edgeToIter)->nodeTo->edgeFrom.erase(edgeFromIter);
				break;
			}
		}
	}

	// find local particle positions
	map<Particle*, Vector3d> particleMap;
	Vector3d position;

	for (set<GraphEdge*>::iterator edgeIter = localGraph->edges.begin(); edgeIter
			!= localGraph->edges.end(); edgeIter++) {
		for (vector<Particle*>::iterator edgeParticleIter =
				(*edgeIter)->particles.begin(); edgeParticleIter
				!= (*edgeIter)->particles.end(); edgeParticleIter++) {
			position = forwardModelFrac(localGraph->nodes.find(
					(*edgeIter)->nodeFrom)->second, (*edgeIter)->odometry,
					(*edgeParticleIter)->fraction);
			if ((*edgeParticleIter)->direction)
				position(2) = constrainAngle(position(2) + pi);

			mexPrintf(
					"Particle %X on edge %X (%X) between %d and %d at fraction %0.4f gives local position %0.4f %0.4f %0.4f\n",
					*edgeParticleIter, *edgeIter, (*edgeParticleIter)->root,
					(*edgeIter)->nodeFrom->frameNum,
					(*edgeIter)->nodeTo->frameNum,
					(*edgeParticleIter)->fraction, position(0), position(1),
					position(2));

			particleMap.insert(pair<Particle*, Vector3d> (*edgeParticleIter,
					position));
		}
	}

	// create equivalent main edge
	mexPrintf("create equivalent main edge\n");

	vector<GraphEdge*> newEdges;

	Vector3d odometry = inverseModel(localGraph->nodes.find(
			node->bestParents.second->nodeTo)->second, localGraph->nodes.find(
			node->bestParents.first->nodeFrom)->second);
	if (getRadius(odometry) < eps) {
		odometry(1) = odometry(1) + 1e-5;
	}
	Matrix3d icov = getInverseCovariance(odometry, Q);

	mexPrintf("Equivalent odometry %0.4 %0.4f %0.4f\n", odometry(0),
			odometry(1), odometry(2));

	GraphEdge* edge = node->bestParents.first->nodeFrom->addEdgeTo(
			node->bestParents.second->nodeTo, odometry, icov);
	mexPrintf("New edge %X from node %X at %d to node %X at %d!\n", edge,
			edge->nodeFrom, edge->nodeFrom->frameNum, edge->nodeTo,
			edge->nodeTo->frameNum);
	edges.push_back(edge);
	newEdges.push_back(edge);

	// add links between from nodes
	mexPrintf("add links between from nodes\n");
	for (vector<GraphEdge*>::iterator edgeFromIter = node->edgeFrom.begin(); edgeFromIter
			!= node->edgeFrom.end(); edgeFromIter++) {
		if ((*edgeFromIter)->nodeFrom != node->bestParents.first->nodeFrom) {
			odometry = inverseModel(localGraph->nodes.find(
					node->bestParents.first->nodeFrom)->second,
					localGraph->nodes.find((*edgeFromIter)->nodeFrom)->second);
			icov = getInverseCovariance(odometry, Q);
			if (getRadius(odometry) < eps) {
				odometry(1) = odometry(1) + 1e-5;
			}

			edge = (*edgeFromIter)->nodeFrom->addEdgeTo(
					node->bestParents.first->nodeFrom, odometry, icov);
			mexPrintf("New edge %X from node %X at %d to node %X at %d!\n", edge,
						edge->nodeFrom, edge->nodeFrom->frameNum, edge->nodeTo,
						edge->nodeTo->frameNum);
			edges.push_back(edge);
			newEdges.push_back(edge);
		}
	}
	// add links between to nodes
	mexPrintf("add links between to nodes\n");
	for (vector<GraphEdge*>::iterator edgeToIter = node->edgeTo.begin(); edgeToIter
			!= node->edgeTo.end(); edgeToIter++) {
		if ((*edgeToIter)->nodeTo != node->bestParents.second->nodeTo) {
			odometry = inverseModel(localGraph->nodes.find(
					node->bestParents.second->nodeTo)->second,
					localGraph->nodes.find((*edgeToIter)->nodeTo)->second);
			icov = getInverseCovariance(odometry, Q);
			if (getRadius(odometry) < eps) {
				odometry(1) = odometry(1) + 1e-5;
			}

			edge = (*edgeToIter)->nodeTo->addEdgeTo(
					node->bestParents.second->nodeTo, odometry, icov);
			mexPrintf("New edge %X from node %X at %d to node %X at %d!\n", edge,
						edge->nodeFrom, edge->nodeFrom->frameNum, edge->nodeTo,
						edge->nodeTo->frameNum);
			edges.push_back(edge);
			newEdges.push_back(edge);
		}
	}

	// update affected particles
	mexPrintf("update affected particles\n");
	Particle* particle;
	Vector3d origin(0,0,0);

	for (map<Particle*, Vector3d>::iterator particleIter = particleMap.begin(); particleIter
			!= particleMap.end(); particleIter++) {
		particle = particleIter->first;
		double likelihood, maxLikelihood = -DBL_MAX;
		double alpha, maxAlpha = 0;
		GraphEdge* maxEdge = newEdges.front();

		for (vector<GraphEdge*>::iterator edgeIter = newEdges.begin(); edgeIter
				!= newEdges.end(); edgeIter++) {
			edge = *edgeIter;
			alpha = get2DForwardDPx_2(
					localGraph->nodes.find(edge->nodeFrom)->second,
					localGraph->nodes.find(edge->nodeTo)->second,
					particleIter->second, edge->icov);
			if (alpha < 0)
				alpha = 0;
			if (alpha > 1)
				alpha = 1;
			likelihood = get2DForwardLogPx(forwardModelFrac(
					localGraph->nodes.find(edge->nodeFrom)->second,
					edge->odometry, alpha), particleIter->second, edge->icov);
			if (likelihood > maxLikelihood) {
				maxLikelihood = likelihood;
				maxAlpha = alpha;
				maxEdge = edge;
			}

		}

		Vector3d maxEdgePos = forwardModelFrac(localGraph->nodes.find(
				maxEdge->nodeFrom)->second, maxEdge->odometry, maxAlpha);

		double forwardAngDiff = angleDistance(maxEdgePos(2),
				particleIter->second(2));
		double reverseAngDiff = angleDistance(
				constrainAngle(maxEdgePos(2) + pi), particleIter->second(2));
		bool maxDirection = abs(forwardAngDiff) > abs(reverseAngDiff);

		particle->root = maxEdge;
		particle->fraction = maxAlpha;
		particle->odoLikelihood = maxLikelihood;
		particle->direction = maxDirection;
		particle->root->particles.push_back(particle);

	}

	// delete node and edges
	mexPrintf("delete node and edges\n");

	for (set<GraphEdge*>::iterator edgeIter = localGraph->edges.begin(); edgeIter
			!= localGraph->edges.end(); edgeIter++) {
		mexPrintf("Deleting edge %X between %X and %X\n",*edgeIter,(*edgeIter)->nodeFrom,(*edgeIter)->nodeTo);
		delete *edgeIter;
	}

	mexPrintf("Deleting local graph\n");
	delete localGraph;

	mexPrintf("Deleting node %X\n",node);
	delete node;

	mexPrintf("Pruning done.\n");

}


void CatGraph::updateParticlePositions(Vector3d odometry,
		unsigned int blankFrames) {
	Particle* particle;
	GraphEdge* edge;
	Vector3d normal;
	Vector3d origin(0, 0, 0);
	Vector3d position;
	Matrix3d icov;

	clearParticleReferences();

	for (vector<Particle*>::iterator particleIter = particles.begin(); particleIter
			!= particles.end(); particleIter++) {
		particle = *particleIter;

		//mexPrintf("Particle %X ",particle);

		// add odometry noise
		normal = Vector3d(randn(0, 1), randn(0, 1), randn(0, 1));
		odometry = odometry + Q * normal;
		position = forwardModel(origin, odometry);

		icov = getInverseCovariance(odometry, Q);
		//mexPrintf("moves %0.4f %0.4f %0.4f\n",odometry(0),odometry(1),odometry(2));

		// build local graph
		double radius = 5 * getRadius(odometry) + 1.0;
		LocalGraph* localGraph = new LocalGraph(particle, radius, frameNum
				- blankFrames);

		//mexPrintf("Local graph has radius %0.4f, %d nodes and %d edges\n",radius, localGraph->nodes.size(),localGraph->edges.size());

		// traverse edges for peaks
		double likelihood, maxLikelihood = -DBL_MAX;
		double alpha, maxAlpha = 0;
		GraphEdge* maxEdge = particle->root;
		bool maxDirection;

		for (set<GraphEdge*>::iterator edgeIter = localGraph->edges.begin(); edgeIter
				!= localGraph->edges.end(); edgeIter++) {
			edge = *edgeIter;

			// forward
			alpha = get2DForwardDPx_2(
					localGraph->nodes.find(edge->nodeFrom)->second,
					localGraph->nodes.find(edge->nodeTo)->second, position,
					icov);
			//mexPrintf("Forward Edge between %d and %d has alpha %0.4f, ",edge->nodeFrom->frameNum,edge->nodeTo->frameNum,alpha);
			if (alpha < 0)
				alpha = 0;
			if (alpha > 1)
				alpha = 1;
			likelihood = get2DForwardLogPx(forwardModelFrac(
					localGraph->nodes.find(edge->nodeFrom)->second,
					edge->odometry, alpha), position, icov);
			//mexPrintf("bound to %0.4f, yields likelihood %0.4f\n",alpha,likelihood);
			if (likelihood > maxLikelihood) {
				maxLikelihood = likelihood;
				maxAlpha = alpha;
				maxEdge = edge;
				//maxDirection = false; // forwards
			}
			/*
			 // reverse
			 alpha = getReverseDPx_2(
			 localGraph->nodes.find(edge->nodeFrom)->second,
			 localGraph->nodes.find(edge->nodeTo)->second, position,
			 icov);
			 mexPrintf("Reverse Edge between %d and %d has alpha %0.4f, ",edge->nodeFrom->frameNum,edge->nodeTo->frameNum,alpha);
			 if (alpha < 0)
			 alpha = 0;
			 if (alpha > 1)
			 alpha = 1;
			 likelihood = getReverseLogPx(forwardModelFrac(localGraph->nodes.find(
			 edge->nodeFrom)->second, edge->odometry, alpha), position,
			 icov);
			 mexPrintf("bound to %0.4f, yields likelihood %0.4f\n",alpha,likelihood);
			 if (likelihood > maxLikelihood) {
			 maxLikelihood = likelihood;
			 maxAlpha = alpha;
			 maxEdge = edge;
			 maxDirection = true; // reverse
			 }
			 */
		}

		Vector3d maxEdgePos = forwardModelFrac(localGraph->nodes.find(
				maxEdge->nodeFrom)->second, maxEdge->odometry, maxAlpha);

		double forwardAngDiff = angleDistance(maxEdgePos(2), position(2));
		double reverseAngDiff = angleDistance(
				constrainAngle(maxEdgePos(2) + pi), position(2));
		maxDirection = abs(forwardAngDiff) > abs(reverseAngDiff);
		//mexPrintf("Forward angle diff %0.4f Reverse angle diff %0.4f\n",forwardAngDiff,reverseAngDiff);

		//mexPrintf("Particle moved %0.4f with likelihood %0.4f\n", getRadius(forwardModelFrac(
		//		localGraph->nodes.find(maxEdge->nodeFrom)->second,
		//		maxEdge->odometry, maxAlpha)),maxLikelihood);

		// update particle position and likelihood
		particle->root = maxEdge;
		particle->fraction = maxAlpha;
		particle->odoLikelihood = maxLikelihood;
		particle->direction = maxDirection;
		particle->root->particles.push_back(particle);

		// delete the local graph
		delete localGraph;
	}

	// update mean odo weight
	// hack for now, set to 0
	odoLikelihood_mean = 0;
}

void CatGraph::updateParticleWeights() {

	// Build map of nodes to compute
	mexPrintf("Building map of nodes\n");
	map<GraphNode*,double> visLikelihoodMap;
	Particle* particle;
	for (vector<Particle*>::iterator particleIter = particles.begin(); particleIter
				!= particles.end(); particleIter++) {
		particle = *particleIter;
		visLikelihoodMap[particle->root->nodeFrom] = -DBL_MAX;
		visLikelihoodMap[particle->root->nodeTo] = -DBL_MAX;
	}
	mexPrintf("Added %d unique nodes to compare\n",visLikelihoodMap.size());

	// Update visual likelihoods for nodes in map
	Bagofwords* currBag = nodes.back()->bagOfWords;
	for (map<GraphNode*, double>::iterator mapIter = visLikelihoodMap.begin(); mapIter
			!= visLikelihoodMap.end(); mapIter++) {
		mapIter->second = fabMap.LLH(*(mapIter->first->bagOfWords),*currBag);
		//BowTemplate currBow;

		//mexPrintf(
		//		"Comparison of node %X at %d to node %X at %d yields %0.4f\n",
		//		nodes.back(), nodes.back()->frameNum, mapIter->first,
		//		mapIter->first->frameNum, mapIter->second);
	}

	double weightSum = 0;
	// Apply likelihoods to particles and update weights
	for (vector<Particle*>::iterator particleIter = particles.begin(); particleIter
					!= particles.end(); particleIter++) {
			particle = *particleIter;

			particle->visLikelihood = (1 - particle->fraction)
				* visLikelihoodMap[particle->root->nodeFrom] + particle->fraction
				* visLikelihoodMap[particle->root->nodeTo];
			particle->weight = particle->weight
					* exp(0.5 * (particle->odoLikelihood + particle->visLikelihood
							/ VIS_DIV));
			// BIG HACK - Multiply by number of words
			//particle->weight *= (double)particle->root->nodeFrom->bagOfWords->getCount()
				+ (double)particle->root->nodeFrom->bagOfWords->getCount();
			weightSum += particle->weight;
	}

	// Generate new place weight
	visLikelihood_mean = z_avg.loglikelihood(*currBag);
	mexPrintf("Average place likelihood yields %0.4f\n",visLikelihood_mean);
	w_mean = NEW_MUL / (double) numParticles * exp(0.5
			* (odoLikelihood_mean + visLikelihood_mean
					/ VIS_DIV));// * 2 * (double)currBag->getSize();
	weightSum += w_mean;

	// Normalise particles
	for (vector<Particle*>::iterator particleIter = particles.begin(); particleIter
			!= particles.end(); particleIter++) {
		particle = *particleIter;
		particle->weight /= weightSum;
		if (particle->weight < eps)
			particle->weight = eps;
	}
	w_mean /= weightSum;

	mexPrintf("New Place likelihood: %%%0.4f\n", w_mean * 100);

}

double CatGraph::calculateESS() {
	// Iterate through particles
	double ess = 0;
	for (vector<Particle*>::iterator particle = particles.begin(); particle
			!= particles.end(); particle++) {
		ess += pow((*particle)->weight * (double) (numParticles + 1) - 1, 2);
	}

	// Add mean particle weight
	ess += pow(w_mean * (double) (numParticles + 1) - 1, 2);
	return 1 / (1 + ess / (double) numParticles);
}

void CatGraph::calculateHypothesis(double hypothesisDist, unsigned int blankFrames) {
	Particle* particle;
	for (vector<Particle*>::iterator particleIter = particles.begin(); particleIter
			!= particles.end(); particleIter++) {
		particle = *particleIter;
		map<Particle*,double> setParticles;
		particle->hypothesis = 0;

		// build local graph
		LocalGraph* localGraph = new LocalGraph(particle, 2 * hypothesisDist,
				frameNum-blankFrames);
		Vector3d position;

		// add particle weights
		for (set<GraphEdge*>::iterator edgeIter = localGraph->edges.begin(); edgeIter
				!= localGraph->edges.end(); edgeIter++) {
			for (vector<Particle*>::iterator edgeParticleIter =
					(*edgeIter)->particles.begin(); edgeParticleIter
					!= (*edgeIter)->particles.end(); edgeParticleIter++) {
				position = forwardModelFrac(localGraph->nodes.find(
						(*edgeIter)->nodeFrom)->second, (*edgeIter)->odometry,
						(*edgeParticleIter)->fraction);
				if ((*edgeParticleIter)->direction)
								position(2) = constrainAngle(position(2)+pi);
				double distance = getRadius(position);
				double forwardAngDiff = angleDistance(position(2), 0);
				double reverseAngDiff = angleDistance(constrainAngle(
						position(2) + pi), 0);
				if (abs(reverseAngDiff) > abs(forwardAngDiff) && distance
						< hypothesisDist) {
					setParticles.insert(pair<Particle*,double>(*edgeParticleIter,distance));
				}
			}
		}

		for (map<Particle*, double>::iterator setParticleIter =
				setParticles.begin(); setParticleIter != setParticles.end(); setParticleIter++) {
			//mexPrintf("Adding %0.4f from Particle %X to Hypothesis %X\n",
			//		setParticleIter->first->weight / (1.0
			//				+ setParticleIter->second), *setParticleIter,
			//		particle);
			particle->hypothesis += setParticleIter->first->weight / (1.0
					+ setParticleIter->second);
		}

		// clear memory
		delete localGraph;
	}

}

void CatGraph::addLoopClosure(Particle* particle) {
	// calculate equivalent odometry

	// big bug - make reverse links work properly
	Vector3d originFrom(0, 0, pi * (double) particle->direction);
	Vector3d originTo(0, 0, 0);
	Vector3d odometry = nodes.back()->edgeFrom.front()->odometry;
	Matrix3d icov;

	Vector3d linkFrom = reverseModel(originFrom, odometry);
	Vector3d linkTo = reverseModelFrac(originTo, particle->root->odometry,
			particle->fraction);

	odometry = inverseModel(linkTo, linkFrom);
	icov = getInverseCovariance(odometry, Q);

	// add new edge
	GraphEdge* edge = nodes.back()->edgeFrom.front()->nodeFrom->addEdgeTo(
			particle->root->nodeFrom, odometry, icov);
	edges.push_back(edge);

	// add information check
	updateList.push_back(particle->root->nodeFrom);

}

void CatGraph::resampleParticles(double spreadFraction, double reverseFraction,
		unsigned int blankFrames) {

	clearParticleReferences();

	// Preallocate vectors for speed
	vector<Particle*> newParticles;
	newParticles.reserve(numParticles);
	vector<double> orderedRand;
	orderedRand.reserve(numParticles);
	vector<double> culmulativeWeight;
	culmulativeWeight.reserve(numParticles);

	// Create culmulative weight and random weight vectors
	double randTotal = 0;
	double weightTotal = 0;
	for (vector<Particle*>::iterator particle = particles.begin(); particle
			!= particles.end(); particle++) {
		culmulativeWeight.push_back(weightTotal + (*particle)->weight * (1
				- spreadFraction));
		weightTotal += (*particle)->weight * (1 - spreadFraction);
		double uniformRand = (double) rand() / RAND_MAX;
		orderedRand.push_back(randTotal + uniformRand);
		randTotal += uniformRand;
	}

	// Resample using Select With Replacement
	unsigned int i = 0;
	unsigned int j = 0;
	while (i < numParticles) {
		if (j < numParticles) {
			if (orderedRand[i] / randTotal < culmulativeWeight[j]) {
				// resample to current particle
				Particle* newParticle = new Particle(*particles[j]);
				double uniformRand = (double) rand() / RAND_MAX;
				if (uniformRand < reverseFraction)
					newParticle->direction = !newParticle->direction;
				newParticle->weight = (1 / (double) numParticles);
				newParticle->root->particles.push_back(newParticle);
				newParticles.push_back(newParticle);
				i++;
			} else {
				j++;
			}
		} else {
			// resample to random location
			Particle* newParticle = randomParticle(blankFrames);
			newParticle->weight = (1 / (double) numParticles);
			//newParticle->root->particles.push_back(newParticle);
			newParticles.push_back(newParticle);
			i++;
		}
	}

	// Swap vectors
	newParticles.swap(particles);
	for (vector<Particle*>::iterator particleIter = newParticles.begin(); particleIter
			!= newParticles.end(); particleIter++) {
		delete *particleIter;
	}
	newParticles.clear();

	mexPrintf("Resampled %d particles\n",particles.size());

}

void CatGraph::clearParticleReferences() {
	for (vector<GraphEdge*>::iterator edge = edges.begin(); edge != edges.end(); edge++) {
		(*edge)->particles.clear();
	}
}

Particle* CatGraph::randomParticle(unsigned int blankFrames) {

	// random numbers go here
	GraphEdge* root = edges.back();
	while (root->frameNum > this->frameNum - blankFrames) {
		root = edges[rand() % (edges.size() - 1)];
	}
	double fraction = (double) rand() / RAND_MAX;
	bool direction = ((double) rand() / RAND_MAX) > 0.5;

	Particle* particle = new Particle(root, fraction, direction);

	return particle;
}

void CatGraph::buildCodebookFromMap(unsigned int numWords, string codebookPath) {

	Mat all_descriptors;
	GraphNode* node;
	mexPrintf("Building training data\n");
	for (vector<GraphNode*>::iterator nodeIter = nodes.begin(); nodeIter
			!= nodes.end(); nodeIter++) {
		node = *nodeIter;
		mexPrintf("Added %d descriptors\n",node->descriptors->rows);
		all_descriptors.push_back(*(node->descriptors));
	}

	cv::TermCriteria tc(cv::TermCriteria::EPS,1000,0.01);
	cv::BOWKMeansTrainer bowTrainer(numWords,tc,1,cv::KMEANS_PP_CENTERS);
	mexPrintf("Clustering descriptors...\n");
	Mat new_codebook = bowTrainer.cluster(all_descriptors);

	mexPrintf("Writing codebook to %s\n",codebookPath.c_str());
	cv::FileStorage fs(codebookPath, cv::FileStorage::WRITE);
	fs << "codebook" << new_codebook;
	fs.release();

	mexPrintf("Releasing data\n");
	all_descriptors.release();
	new_codebook.release();

}

void CatGraph::buildChowLiuFromMap(double infoThreshold, string chowLiuPath) {

	TrainData trainData;
	trainData.num_samples = nodes.size();
	trainData.sample_size = codebook.rows;

	trainData.data = new bool*[trainData.num_samples];
	for(int i = 0; i < trainData.num_samples; i++) {
		trainData.data[i] = new bool[trainData.sample_size];
		memset(trainData.data[i],0,trainData.sample_size*sizeof(bool));
	}

	mexPrintf("Build training data matrix with %d rows of %d features\n",
			trainData.num_samples, trainData.sample_size);

	unsigned int nodeIndex = 0;
	GraphNode* node;
	for (vector<GraphNode*>::iterator nodeIter = nodes.begin(); nodeIter
				!= nodes.end(); nodeIter++) {
		node = *nodeIter;
		vector<cv::DMatch> matches;
		matcher->match(*(node->descriptors),codebook,matches);
		for (vector<cv::DMatch>::iterator matchIter = matches.begin(); matchIter
				!= matches.end(); matchIter++) {
			trainData.data[nodeIndex][matchIter->trainIdx] = true;
		}
		nodeIndex++;
		mexPrintf("Extracted %d of %d features\n",nodeIndex,trainData.num_samples);
	}
	mexPrintf("Making absolutes\n");
	trainData.make_absolutes();

	mexPrintf("Building Chow-Liu tree:\n");
	ClTree new_clTree;
	new_clTree.make(&trainData, infoThreshold);

	mexPrintf("Saving Chow-Liu tree to %s\n",chowLiuPath.c_str());
	new_clTree.save(chowLiuPath);

}

void CatGraph::recomputeWords() {

	// Memory leaky...
	GraphNode* node;
	for (vector<GraphNode*>::iterator nodeIter = nodes.begin(); nodeIter
			!= nodes.end(); nodeIter++) {
		node = *nodeIter;

		mexPrintf("Loading descriptors from file\n");
		cv::FileStorage descriptorFile(node->descriptorPath,cv::FileStorage::READ);
		cv::Mat* descriptors = new cv::Mat;

		descriptorFile["descriptors"] >> (*descriptors);
		descriptorFile.release();

		vector<cv::DMatch> matches;

		mexPrintf("Creating bag of words\n");
		node->bagOfWords = new Bagofwords(codebook.rows);
		matcher->match(*descriptors, codebook, matches);
		for (vector<cv::DMatch>::iterator matchIter = matches.begin(); matchIter
				!= matches.end(); matchIter++) {
			node->bagOfWords->add(matchIter->trainIdx);
		}

		mexPrintf("Node %X with %d %d-D descriptors has %d unique words\n",
				node, descriptors->rows, descriptors->cols,
				node->bagOfWords->getCount());

		descriptors->release();
	}
}

void CatGraph::recomputeLikelihoods() {
	GraphNode* node;
	Bagofwords* currBag = nodes.back()->bagOfWords;
	for (vector<GraphNode*>::iterator nodeIter = nodes.begin(); nodeIter
				!= nodes.end(); nodeIter++) {
		node = *nodeIter;
		// Backwards?
		node->visLikelihood = fabMap.LLH(*(node->bagOfWords),*currBag);
	}
	visLikelihood_mean = z_avg.loglikelihood(*currBag);
}

void CatGraph::recomputeInformation() {
	for (vector<GraphNode*>::iterator nodeIter = nodes.begin(); nodeIter
					!= nodes.end(); nodeIter++) {
		updateList.push_back(*nodeIter);
	}
}

// default constructors

CatGraph::CatGraph() {
	nodes.clear();
	edges.clear();
	particles.clear();
}

CatGraph::~CatGraph() {
	nodes.clear();
	edges.clear();
	particles.clear();
}

CatGraph::CatGraph(const CatGraph & other) {
	frameNum = other.frameNum;
	numParticles = other.numParticles;
	numNodes = other.numNodes;
	Q = other.Q;
	u_mean = other.u_mean;
	nodes = other.nodes;
	edges = other.edges;
	particles = other.particles;
}

CatGraph & CatGraph::operator=(const CatGraph & other) {
	frameNum = other.frameNum;
	numParticles = other.numParticles;
	numNodes = other.numNodes;
	Q = other.Q;
	u_mean = other.u_mean;
	nodes = other.nodes;
	edges = other.edges;
	particles = other.particles;
	return *this;
}

// serialize and deserialize
void CatGraph::serializeGraph(string fileName) {
	ofstream ofs(fileName.c_str());
	boost::archive::text_oarchive oa(ofs);
	//boost::archive::xml_oarchive oa(ofs);
	//oa << *this;
	oa << BOOST_SERIALIZATION_NVP(*this);
}

void CatGraph::deserializeGraph(string fileName) {
	ifstream ifs(fileName.c_str());
	boost::archive::text_iarchive ia(ifs);
	//boost::archive::xml_iarchive ia(ifs);
	//ia >> *this;
	mexPrintf("Deserializing into CatGraph at %X\n",this);
	ia >> BOOST_SERIALIZATION_NVP(*this);
	mexPrintf("--- NOTE: Call initialiseVision and then recomputeWords first ---\n");
}
